#include "common.h"


// 1. RDMA连接相关的资源
static struct ibv_pd *pd      = NULL;
static struct ibv_cq *clnt_cq = NULL;
static struct ibv_qp *clnt_qp;


static struct rdma_cm_id *cm_clnt_id = NULL;   // 客户端CM ID，类型与Socket的客户端套接字
static struct ibv_qp_init_attr qp_attr;        // 创建QP的属性信息


// CM事件管道
// RDMA IO通信事件管道
static struct rdma_event_channel *cm_event_channel      = NULL;
static struct ibv_comp_channel   *io_comp_channel = NULL;


// RDMA操作的真实数据和元数据的MR信息
static struct ibv_mr *client_meta_mr = NULL, 
					 *client_src_mr  = NULL, 
					 *client_dst_mr  = NULL, 
					 *server_meta_mr = NULL;

// 服务端可客户端的元数据信息
// 在真正进行RDMA操作前使用Post Send/Recv交换该元信息
static struct meta_attr client_meta_attr, server_meta_attr;


// 接收和发送的WR和SGE
static struct ibv_send_wr client_send_wr, *bad_client_send_wr = NULL;
static struct ibv_recv_wr server_recv_wr, *bad_server_recv_wr = NULL;
static struct ibv_sge client_send_sge, server_recv_sge;



// 客户端字符串保存的地址，客户端通过RDMA WRITE将该地址的数据发送给服务端
// 客户端通过RDMA READ将服务端的数据读取到该地址
static char *src = NULL, *dst = NULL; 



//! 1. 准备RDMA操作相关的资源
static int setup(struct sockaddr_in *s_addr)
{
	int ret = -1;
	struct rdma_cm_event *cm_event = NULL;


	//! 1. 创建CM(连接管理器)事件管道
	cm_event_channel = rdma_create_event_channel();
	if (!cm_event_channel) {
		rdma_error("创建CM事件管道失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("创建CM事件管道成功: %p \n", cm_event_channel);


	//! 2. 创建CM(类似于以太方式的套接字)
	ret = rdma_create_id(cm_event_channel, 
			&cm_clnt_id, NULL, RDMA_PS_TCP);
	if (ret) {
		rdma_error("CM创建失败, 错误码: %d \n", -errno); 
		return -errno;
	}


	//! 3. 解析服务端地址
	//! 3.1 rdma_resolve_addr主要解析联通性(异步)
	ret = rdma_resolve_addr(cm_clnt_id, NULL,         // 客户端地址留空自动解析
							(struct sockaddr*)s_addr, // 解析服务端地址
							2000                      // 解析超时时间
						);
	if (ret) {
		rdma_error("解析服务端地址失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("等待CM事件: RDMA_CM_EVENT_ADDR_RESOLVED\n");

	// 解析时间阻塞回调处理
	ret  = cm_event_callback(cm_event_channel, 
			RDMA_CM_EVENT_ADDR_RESOLVED, &cm_event);
	if (ret) {
		rdma_error("服务端地址解析失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("服务端地址解析CM时间释放失败, 错误码: %d\n", -errno);
		return -errno;
	}
	rdma_debug("服务端RDMA地址解析成功...\n");


	//! 3.2 解析通往Server端的RDMA设备的完整路径
	ret = rdma_resolve_route(cm_clnt_id, 2000);
	if (ret) {
		rdma_error("解析服务端路由解析失败, 错误码: %d \n", -errno);
	       return -errno;
	}
	rdma_debug("等待CM事件: RDMA_CM_EVENT_ROUTE_RESOLVED\n");
	ret = cm_event_callback(cm_event_channel, 
			RDMA_CM_EVENT_ROUTE_RESOLVED,&cm_event);
	if (ret) {
		rdma_error("从CM管道接收路由解析失败, 错误码: %d \n", ret);
		return ret;
	}

	
	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("释放服务端RDMA路由CM事件失败, 错误码: %d \n", -errno);
		return -errno;
	}
	printf("[INFO] 尝试连接到RDMA服务: %s port: %d \n", 
		inet_ntoa(s_addr->sin_addr), ntohs(s_addr->sin_port));

	

	//! 4. 创建RDMA通信需要的资源
	//! 4.1 创建PD
	pd = ibv_alloc_pd(cm_clnt_id->verbs);
	if (!pd) {
		rdma_error("客户端PD创建失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("客户端PD创建成功: %p \n", pd);



	//! 4.2 创建IO完成管道
	io_comp_channel = ibv_create_comp_channel(cm_clnt_id->verbs);
	if (!io_comp_channel) {
		rdma_error("创建IO完成事件管道失败, 错误码: %d\n",
			       -errno);
	return -errno;
	}
	rdma_debug("创建IO完成事件管道成功: %p \n", io_comp_channel);



	//! 4.3 创建完成队列CQ
	clnt_cq = ibv_create_cq(cm_clnt_id->verbs, // RDMA设备上下文 
							CQ_CAPACITY,       // 完成队列的荣阿玲 
							NULL,              // 用户上下文
							io_comp_channel,   // IO完成事件管道
							0                  // 信号向量
						);
	if (!clnt_cq) {
		rdma_error("CQ创建失败, 错误码: %d \n", -errno);
		return -errno;
	}
	rdma_debug("创建CQ成功: %p 持有 %d 个元素! \n", clnt_cq, clnt_cq->cqe);


	// int ibv_req_notify_cq(struct ibv_cq *cq, 
	//                      int solicited_only);   // 通知出发条件的标志
	//                      0: 通知所有事件
	//                      1: 只有被标记为"solicited"的工作完成才会触发通知
	ret = ibv_req_notify_cq(clnt_cq, 0);
	if (ret) {
		rdma_error("CQ通知条件设置失败, 错误码: %d\n", -errno);
		return -errno;
	}

	//! 4.4 创建QP
	//! 4.4.1 初始化QP属性
	bzero(&qp_attr, sizeof qp_attr);
	qp_attr.cap.max_recv_sge = MAX_SGE;    // 接收SGE容量
	qp_attr.cap.max_recv_wr  = MAX_WR;     // 接收WR容量
	qp_attr.cap.max_send_sge = MAX_SGE;    // 发送SGE容量
	qp_attr.cap.max_send_wr  = MAX_WR;     // 发送WR容量
	qp_attr.qp_type          = IBV_QPT_RC; // QP类型


	// 这里我们使用相同的完成队列, 也可以选择使用不同的队列
	qp_attr.recv_cq = clnt_cq; // Recv事件完成队列
	qp_attr.send_cq = clnt_cq; // Send事件完成队列
    
	//! 4.4.2 创建QP
	ret = rdma_create_qp(cm_clnt_id,    // 指定QP的连接
		       pd,                      // 指定QP的保护域
		       &qp_attr);               // 指定QP的属性
	if (ret) {
		rdma_error("QP创建失败, 错误码: %d \n", -errno);
	    return -errno;
	}


	clnt_qp = cm_clnt_id->qp;
	rdma_debug("QP创建成功: %p \n", clnt_qp);

	return 0;
}



//! 2. 发送接收服务端元数据的Post Recv请求
static int post_recv_meta()
{
	int ret = -1;
	//! 1. 注册服务端元数据MR(用于存放服务端元数据)
	server_meta_mr = mr_register(pd, &server_meta_attr,
							sizeof(server_meta_attr),
							(IBV_ACCESS_LOCAL_WRITE));
	if(!server_meta_mr){
		rdma_error("创建服务端元数据MR失败, -ENOMEM\n");
		return -ENOMEM;
	}

	//! 2. 设置服务端元数据信息
	server_recv_sge.addr   = (uint64_t) server_meta_mr->addr;
	server_recv_sge.length = (uint32_t) server_meta_mr->length;
	server_recv_sge.lkey   = (uint32_t) server_meta_mr->lkey;
	
	/* now we link it to the request */

	//! 3. 创建服务端元数据WR
	bzero(&server_recv_wr, sizeof(server_recv_wr));
	server_recv_wr.sg_list = &server_recv_sge;
	server_recv_wr.num_sge = 1;

	//! 4. 向本地硬件下发Post Recv请求, 
	//   请求服务端元数据, 存储到server_meta_mr区域
	ret = ibv_post_recv(clnt_qp ,
		      &server_recv_wr,
		      &bad_server_recv_wr);
	if (ret) {
		rdma_error("向本地硬件下发Post Recv请求失败, 错误码: %d \n", ret);
		return ret;
	}

	rdma_debug("Receive buffer pre-posting is successful \n");
	return 0;
}



//! 3. 连接到服务端
static int conn_to_server() 
{
	int ret = -1;
	struct rdma_conn_param conn_param;
	struct rdma_cm_event *cm_event = NULL;

	// 1. 创建链接服务器的参数
	bzero(&conn_param, sizeof(conn_param));
	conn_param.initiator_depth     = 3;
	conn_param.responder_resources = 3;
	conn_param.retry_count         = 3; // 链接重试次数

	
	// 2. 异步连接到服务端
	// 2.1 发送链接请求
	ret = rdma_connect(cm_clnt_id, &conn_param);
	if (ret) {
		rdma_error("向服务器发送链接请求失败, 错误码: %d\n", -errno);
		return -errno;
	}
	rdma_debug("等待CM事件: RDMA_CM_EVENT_ESTABLISHED\n");

	// 2.2 阻塞等待服务端连接建立事件
	ret = cm_event_callback(cm_event_channel, 
			RDMA_CM_EVENT_ESTABLISHED,
			&cm_event);
	if (ret) {
		rdma_error("获取服务端连接建立CM事件失败, 错误码: %d \n", ret);
	       return ret;
	}

	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("释放服务端连接建立CM事件失败, 错误码: %d\n", -errno);
		return -errno;
	}
	printf("[INFO ] 客户端连接服务端成功... \n");
	return 0;
}




//! 4. 与服务端交换元数据
static int xchange_meta()
{
	int ret = -1;
	struct ibv_wc wc[2];
	// 1. 准备客户端数据信息MR
	client_src_mr = mr_register(pd, src, strlen(src),
							(IBV_ACCESS_LOCAL_WRITE|
							IBV_ACCESS_REMOTE_READ|
							IBV_ACCESS_REMOTE_WRITE));
	if(!client_src_mr){
		rdma_error("注册客户端数据地址MR失败, 错误码: %d \n", ret);
		return ret;
	}

	// 2. 准备客户端元数据信息
	// 2.1 准备客户端元数据
	client_meta_attr.address  = (uint64_t)client_src_mr->addr; 
	client_meta_attr.length   = (uint32_t)client_src_mr->length; 
	client_meta_attr.key.lkey = (uint32_t)client_src_mr->lkey;
	
	// 2.2 注册客户端元数据MR
	client_meta_mr = mr_register(pd,
			&client_meta_attr,
			sizeof(client_meta_attr),
			IBV_ACCESS_LOCAL_WRITE);
	if(!client_meta_mr) {
		rdma_error("注册客户端元数据地址MR失败, 错误码: %d %d \n", ret);
		return ret;
	}


	// 3. 准备客户端发送SGE
	client_send_sge.addr   = (uint64_t)client_meta_mr->addr;
	client_send_sge.length = (uint32_t)client_meta_mr->length;
	client_send_sge.lkey   = (uint32_t)client_meta_mr->lkey;

	// 4. 准备发送本机元数据WR
	bzero(&client_send_wr, sizeof(client_send_wr));
	client_send_wr.sg_list    = &client_send_sge;
	client_send_wr.num_sge    = 1;
	client_send_wr.opcode     = IBV_WR_SEND;
	client_send_wr.send_flags = IBV_SEND_SIGNALED;
	
	// 5. 提交客户端元数据WR
	ret = ibv_post_send(clnt_qp, &client_send_wr, &bad_client_send_wr);
	if (ret) {
		rdma_error("客户端元数据Post Send任务下发失败, 错误码: %d \n", -errno);
		return -errno;
	}


	// 6. 处理前面下发的接收服务端的Post Recv和刚已提交的发送本机元数据的Post Send请求
	ret = wc_event_callback(io_comp_channel, wc, 2);
	if(ret != 2) {
		rdma_error("未成功获取2个WC, 错误码: %d \n",
				ret);
		return ret;
	}
	rdma_debug("本机向服务端发送元数据成功，并成功获取服务端元数据和访问凭证: \n");

	show_meta(&server_meta_attr);
	return 0;
}



//! 5. RDMA WRITE READ操作
// 5.1 将src地址的数据通过RDMA WRITE写入到服务端	
// 5.2 通过RDMA READ操作将远端数据写入本地dst地址
static int rdma_write_read() 
{
	int ret = -1;
	struct ibv_wc wc;

	// 1. 准备用于RDMA READ操作的MR
	client_dst_mr = mr_register(pd, dst, strlen(src),
						(IBV_ACCESS_LOCAL_WRITE | 
						IBV_ACCESS_REMOTE_WRITE | 
						IBV_ACCESS_REMOTE_READ));
	if (!client_dst_mr) {
		rdma_error("创建客户端RDMA READ操作MR失败, -ENOMEM\n");
		return -ENOMEM;
	}



	// 2. 准备客户端RDMA WRITE操作的SGE
	client_send_sge.addr   = (uint64_t)client_src_mr->addr;
	client_send_sge.length = (uint32_t)client_src_mr->length;
	client_send_sge.lkey   = (uint32_t)client_src_mr->lkey;


	// 3. 准备客户端RDMA WRITE操作的WR
	bzero(&client_send_wr, sizeof(client_send_wr));
	client_send_wr.sg_list    = &client_send_sge;
	client_send_wr.num_sge    = 1;
	client_send_wr.opcode     = IBV_WR_RDMA_WRITE;   //! 注意opcode
	client_send_wr.send_flags = IBV_SEND_SIGNALED;

	// 配置远程访问的地址和密钥
	client_send_wr.wr.rdma.rkey = server_meta_attr.key.rkey;
	client_send_wr.wr.rdma.remote_addr = server_meta_attr.address;

	// 4. 向服务端发送RDMA WRITE请求
	ret = ibv_post_send(clnt_qp, &client_send_wr, &bad_client_send_wr);
	if (ret) {
		rdma_error("客户端RDMA WRITE Post Send任务下发失败, 错误码: %d \n", -errno);
		return -errno;
	}

	// 5. 阻塞等外WC事件
	ret = wc_event_callback(io_comp_channel, &wc, 1);
	if(ret != 1) {
		rdma_error("获取服务端RDMA WRITE的WC失败, 错误码: %d \n",ret);
		return ret;
	}
	rdma_debug("客户端的RDMA WRITE操作成功! \n");


	///////////////////////////////////////////////////////////////
	// 现从远程Server端通过RDMA READ操作读取会数据到dst内存空间
	
	client_send_sge.addr   = (uint64_t)client_dst_mr->addr;
	client_send_sge.length = (uint32_t)client_dst_mr->length;
	client_send_sge.lkey   = (uint32_t)client_dst_mr->lkey;


	bzero(&client_send_wr, sizeof(client_send_wr));
	client_send_wr.sg_list    = &client_send_sge;
	client_send_wr.num_sge    = 1;
	client_send_wr.opcode     = IBV_WR_RDMA_READ;     //! 注意opcode
	client_send_wr.send_flags = IBV_SEND_SIGNALED;
	
	// 配置远程访问的地址和密钥
	client_send_wr.wr.rdma.rkey = server_meta_attr.key.rkey;
	client_send_wr.wr.rdma.remote_addr = server_meta_attr.address;

	//  向服务端发送RDMA READ请求
	ret = ibv_post_send(clnt_qp, &client_send_wr, &bad_client_send_wr);
	if (ret) {
		rdma_error("客户端RDMA READ Post Send任务下发失败, 错误码: : %d \n", 
				-errno);
		return -errno;
	}


	ret = wc_event_callback(io_comp_channel, 
			&wc, 1);
	if(ret != 1) {
		rdma_error("获取服务端RDMA READ的WC失败, 错误码: %d \n",
				ret);
		return ret;
	}
	rdma_debug("客户端的RDMA READ操作成功! \n");

	return 0;
}



//! 6. 验证dst与src内存的一致性
static int res_comp() 
{
	return memcmp((void*) src, (void*) dst, strlen(src));
}



//! 7. 清理RDMA相关的资源
static int cleanup()
{
	int ret = -1;
	struct rdma_cm_event *cm_event = NULL;

	// 1. 创建断开连接的请求
	ret = rdma_disconnect(cm_clnt_id);
	if (ret) {
		rdma_error("客户端创建连接断开请求失败, 错误码 %d \n", -errno);
	}
	ret = cm_event_callback(cm_event_channel, 
			RDMA_CM_EVENT_DISCONNECTED,
			&cm_event);
	if (ret) {
		rdma_error("客户端获取RDMA_CM_EVENT_DISCONNECTED时间失败, 错误码: %d\n", ret);
	}
	ret = rdma_ack_cm_event(cm_event);
	if (ret) {
		rdma_error("释放断开连接CM事件失败, 错误码: %d\n", -errno);
	}

	// 2. 销毁RDMA通信相关的资源
	// 2.1 销毁QP
	rdma_destroy_qp(cm_clnt_id);
	ret = rdma_destroy_id(cm_clnt_id);
	if (ret) {
		rdma_error("释放客户端ID失败, 错误码: %d \n", -errno);
	}


	// 2.2 销毁CQ
	ret = ibv_destroy_cq(clnt_cq);
	if (ret) {
		rdma_error("释放完成队列失败, 错误码: %d \n", -errno);
	}

	// 2.3 销毁完成通道
	ret = ibv_destroy_comp_channel(io_comp_channel);
	if (ret) {
		rdma_error("释放IO完成通道失败, 错误码: %d \n", -errno);
	}


	// 2.4 反注册MR
	mr_deregister(server_meta_mr);
	mr_deregister(client_meta_mr);	
	mr_deregister(client_src_mr);	
	mr_deregister(client_dst_mr);	


	free(src);
	free(dst);

	// 2.5 销毁PD
	ret = ibv_dealloc_pd(pd);
	if (ret) {
		rdma_error("释放PD失败, 错误码: %d \n", -errno);
		// we continue anyways;
	}

	rdma_destroy_event_channel(cm_event_channel);
	printf("[INFO]客户端RDMA资源释放成功... \n");

	return 0;
}


//! 打印帮助信息
void usage() {
	printf("Usage:\n");
	printf("rdma_client: [-a <server_addr>] [-p <server_port>] -s string (required)\n");
	printf("(default IP is 127.0.0.1 and port is %d)\n", RDMA_PORT);
	exit(1);
}



int main(int argc, char **argv) {
	int ret, option;
	struct sockaddr_in server_sock;
	bzero(&server_sock, sizeof server_sock);

	server_sock.sin_family      = AF_INET;
	server_sock.sin_addr.s_addr = htonl(INADDR_LOOPBACK);


	src = dst = NULL; 
	
	// 解析命行参数
	while ((option = getopt(argc, argv, "s:a:p:")) != -1) {
		switch (option) {
			case 's':    //! 1. string: 客户端WRITE给服务端的内容
				printf("传入字符串: %s, 字节长度 %u \n", optarg, 
						(unsigned int) strlen(optarg));
				src = calloc(strlen(optarg) , 1);
				if (!src) {
					rdma_error("传入字符串内存分配失败: -ENOMEM\n");
					return -ENOMEM;
				}
				
				// 命令行参数传入的字符串复制到src
				strncpy(src, optarg, strlen(optarg));
				dst = calloc(strlen(optarg), 1);
				if (!dst) {
					rdma_error("Failed to allocate destination memory, -ENOMEM\n");
					free(src);
					return -ENOMEM;
				}
				break;
			case 'a':    //! 2. 服务端RDMA网卡IP 
				// 通过命令行参数传入的IP地址复制到server_sock
				ret = get_addr(optarg, (struct sockaddr*) &server_sock);
				if (ret) {
					rdma_error("Invalid IP \n");
					return ret;
				}
				break;
			case 'p':    //! 3. 服务端RDMA端口
				server_sock.sin_port = htons(strtol(optarg, NULL, 0)); 
				break;
			default:
				usage();
				break;
			}
		}

		// 如果为传入服务端端口, 使用默认值
		if (!server_sock.sin_port) {
			server_sock.sin_port = htons(RDMA_PORT);
		}
		if (src == NULL) {
			printf("Please provide a string to copy \n");
			usage();
  	}

	ret = setup(&server_sock);
	if (ret) { 
		rdma_error("准备客户端RDMA资源失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = post_recv_meta(); 
	if (ret) { 
		rdma_error("提交请求服务端元数据的Post Recv请求失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = conn_to_server();
	if (ret) { 
		rdma_error("链接到服务端失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = xchange_meta();
	if (ret) {
		rdma_error("与服务端叫喊交换元数据失败, 错误码: %d \n", ret);
		return ret;
	}

	ret = rdma_write_read();
	if (ret) {
		rdma_error("RDMA WRITE READ操作失败, 错误码: %d \n", ret);
		return ret;
	}

	if (res_comp()) {
		rdma_error("失败:src与dst数据不匹配! \n");
	} else {
		printf("...\n成功: src与dst数据匹配!\n");
	}

	ret = cleanup();
	if (ret) {
		rdma_error("清理客户端RDMA资源失败, 错误码: %d \n", ret);
	}
	
	return ret;
}

